'''
@author: kohill
@Date: 2018-8-7
This file is trying to implement a dataset wrapper for helmet detection task,
also, some other useful functions are included, like coco converter (as it's a subclass of BBOX_Base_Dataset).
'''
import os, logging
import numpy as np
import mxnet as mx
import xml.etree.ElementTree as ET
import cv2
from lib.data.bbox_dataset import BBOX_Base_Dataset


class HelmetDataset(BBOX_Base_Dataset):
    def parser_xml(self, xml_path):
        oneimg = {}
        oneimg['bboxes'] = []
        try:
            dom = ET.parse(xml_path)
        except Exception as e:
            logging.warning("{}_{}".format(e, xml_path))
            return None
        root = dom.getroot()
        img_name = root.findall('filename')[0].text
        oneimg['img_path'] = xml_path[:-4] + ".jpg"
        oneimg['file_name'] = img_name
        for objects in root.findall('object'):
            name = objects.find('name').text
            points = list(objects.find('bndbox'))
            if len(points) != 4:
                logging.warning("find illegal label in file:{}.xml. ".format(img_name))
                print(points)
                return None
            xmin = int(points[0].text)
            ymin = int(points[1].text)
            xmax = int(points[2].text)
            ymax = int(points[3].text)

            if str(name).startswith("head"):  # type: str
                oneimg['bboxes'].append([xmin, ymin, xmax, ymax, 0, 0])
        return oneimg

    def __init__(self,
                 XML_ROOT,
                 transform=None,
                 *args, **kwargs
                 ):
        super(HelmetDataset, self).__init__(*args, **kwargs)
        xml_path = XML_ROOT

        csv_path = "/data1/zyx/yks/dataset/helmet_his/train_labels.csv"
        helmet_bbox = {}
        for l in open(csv_path):
            try:
                name, bbox = l.strip().split(',')
                x0, y0, x1, y1 = bbox.strip().split(' ')
                x0 = int(x0)
                y0 = int(y0)
                x1 = int(x1)
                y1 = int(y1)
                try:
                    helmet_bbox[name].append([x0, y0, x1, y1, 1, 0])
                except Exception as e:
                    helmet_bbox[name] = [[x0, y0, x1, y1, 1, 0]]
            except ValueError as e:
                pass
        self.obj = []
        for x, _, z in os.walk(xml_path):
            for name in z:
                if str(name).endswith(".xml"):
                    sxml_path = os.path.join(x, name)
                    oneimg = self.parser_xml(sxml_path)
                    if oneimg is not None and len(oneimg["bboxes"]) > 0:
                        self.obj.append(oneimg)
                        try:
                            oneimg["bboxes"] += helmet_bbox[oneimg["file_name"] + ".jpg"]
                        except KeyError as e:
                            pass
        self.classes = ["head", "helmet"]
        self._transform = transform

    def __len__(self):
        return len(self.obj)

    def __getitem__(self, idx):
        oneimg = self.obj[idx]

        img_path = oneimg["img_path"]
        assert os.path.exists(img_path), img_path

        img_ori = cv2.imread(img_path)[:, :, ::-1]
        bboxes = oneimg['bboxes']
        bboxes = np.array(bboxes)
        if self._transform is not None:
            img_ori, bboxes = self._transform(img_ori, bboxes)
        return mx.nd.array(img_ori), bboxes

    def at_with_image_path(self, idx):
        oneimg = self.obj[idx]
        img_path = oneimg["img_path"]
        bboxes = oneimg['bboxes']
        bboxes = np.array(bboxes)
        return img_path, bboxes


def getDataset(XML_ROOT=None):
    all_dataset = HelmetDataset(XML_ROOT=XML_ROOT)
    return all_dataset


if __name__ == '__main__':
    train_dataset = getDataset(XML_ROOT="/data1/zyx/yks/dataset/helmet/bulk/train")
    # train_dataset.to_roidb("/data1/zyx/yks/dataset/helmet//train.roidb")
    train_dataset.viz()
